/*
 * Copyright (c) 2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
interface ArkTools {
    long timeInUs(Object arg);
}

class NBodyConst {
    public static final double PI = Math.PI;
    public static final double SOLAR_MASS = 4 * NBodyConst.PI * NBodyConst.PI;
    public static final double DAYS_PER_YEAR = 365.24;
}

class Body {
    public double x;
    public double y;
    public double z;
    public double vx;
    public double vy;
    public double vz;
    public double mass;

    Body(double x, double y, double z, double vx, double vy, double vz, double mass) {
        this.x = x;
        this.y = y;
        this.z = z;
        this.vx = vx;
        this.vy = vy;
        this.vz = vz;
        this.mass = mass;
    }

    public static Body jupiter() {
        return new Body(
                4.84143144246472090e+00,
                -1.16032004402742839e+00,
                -1.03622044471123109e-01,
                1.66007664274403694e-03 * NBodyConst.DAYS_PER_YEAR,
                7.69901118419740425e-03 * NBodyConst.DAYS_PER_YEAR,
                -6.90460016972063023e-05 * NBodyConst.DAYS_PER_YEAR,
                9.54791938424326609e-04 * NBodyConst.SOLAR_MASS);
    }

    public static Body saturn() {
        return new Body(
                8.34336671824457987e+00,
                4.12479856412430479e+00,
                -4.03523417114321381e-01,
                -2.76742510726862411e-03 * NBodyConst.DAYS_PER_YEAR,
                4.99852801234917238e-03 * NBodyConst.DAYS_PER_YEAR,
                2.30417297573763929e-05 * NBodyConst.DAYS_PER_YEAR,
                2.85885980666130812e-04 * NBodyConst.SOLAR_MASS);
    }

    public static Body uranus() {
        return new Body(
                1.28943695621391310e+01,
                -1.51111514016986312e+01,
                -2.23307578892655734e-01,
                2.96460137564761618e-03 * NBodyConst.DAYS_PER_YEAR,
                2.37847173959480950e-03 * NBodyConst.DAYS_PER_YEAR,
                -2.96589568540237556e-05 * NBodyConst.DAYS_PER_YEAR,
                4.36624404335156298e-05 * NBodyConst.SOLAR_MASS);
    }

    public static Body neptune() {
        return new Body(
                1.53796971148509165e+01,
                -2.59193146099879641e+01,
                1.79258772950371181e-01,
                2.68067772490389322e-03 * NBodyConst.DAYS_PER_YEAR,
                1.62824170038242295e-03 * NBodyConst.DAYS_PER_YEAR,
                -9.51592254519715870e-05 * NBodyConst.DAYS_PER_YEAR,
                5.15138902046611451e-05 * NBodyConst.SOLAR_MASS);
    }

    public static Body sun() {
        return new Body(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, NBodyConst.SOLAR_MASS);
    }

    public static final Body[] bodies = { sun(), jupiter(), saturn(), uranus(),
        neptune() };
}

class NBodyHandler {
    public static void offsetMomentum(Body[] bodies) {
        double px = 0.0;
        double py = 0.0;
        double pz = 0.0;
        int size = bodies.length;
        for (int i = 0; i < size; i++) {
            Body body = bodies[i];
            double mass = body.mass;
            px += body.vx * mass;
            py += body.vy * mass;
            pz += body.vz * mass;
        }

        Body body = bodies[0];
        body.vx = -px / NBodyConst.SOLAR_MASS;
        body.vy = -py / NBodyConst.SOLAR_MASS;
        body.vz = -pz / NBodyConst.SOLAR_MASS;

        // System.out.println("exec offsetMomentum: body " + " x = " + body.x + " y = " + body.y + " z = " + body.z + " vx = " + body.vx + " vy = " + body.vy + " vz = " + body.vz + " mass = " + body.mass);
    }

    public static void advance(Body[] bodies, double dt) {
        int size = bodies.length;
        for (int i = 0; i < size; i++) {
            Body bodyi = bodies[i];
            double vxi = bodyi.vx;
            double vyi = bodyi.vy;
            double vzi = bodyi.vz;
            double xi = bodyi.x;
            double yi = bodyi.y;
            double zi = bodyi.z;
            double massi = bodyi.mass;

            for (int j = i + 1; j < size; j++) {
                Body bodyj = bodies[j];
                double dx = xi - bodyj.x;
                double dy = yi - bodyj.y;
                double dz = zi - bodyj.z;

                double d2 = dx * dx + dy * dy + dz * dz;
                double mag = dt / (d2 * Math.sqrt(d2));
                double massiMag = massi * mag;

                double massj = bodyj.mass;
                double massjMag = massj * mag;
                vxi -= dx * massjMag;
                vyi -= dy * massjMag;
                vzi -= dz * massjMag;

                bodyj.vx += dx * massiMag;
                bodyj.vy += dy * massiMag;
                bodyj.vz += dz * massiMag;
            }
            bodyi.vx = vxi;
            bodyi.vy = vyi;
            bodyi.vz = vzi;

            bodyi.x += dt * vxi;
            bodyi.y += dt * vyi;
            bodyi.z += dt * vzi;

            // System.out.println("exec advance: body[" + i + "]" + " x = " + bodyi.x + " y = " + bodyi.y + " z = " + bodyi.z + " vx = " + bodyi.vx + " vy = " + bodyi.vy + " vz = " + bodyi.vz + " mass = " + bodyi.mass);
        }
    }

    public static double energy(Body[] bodies) {
        double e = 0.0;
        int size = bodies.length;

        for (int i = 0; i < size; i++) {
            Body bodyi = bodies[i];

            e += 0.5 * bodyi.mass * (bodyi.vx * bodyi.vx + bodyi.vy * bodyi.vy + bodyi.vz * bodyi.vz);

            for (int j = i + 1; j < size; j++) {
                Body bodyj = bodies[j];
                double dx = bodyi.x - bodyj.x;
                double dy = bodyi.y - bodyj.y;
                double dz = bodyi.z - bodyj.z;

                double distance = Math.sqrt(dx * dx + dy * dy + dz * dz);
                e -= (bodyi.mass * bodyj.mass) / distance;
            }
        }

        // System.out.println("exec energy: result e = " + e);

        return e;
    }
}

/*
 * @State
 * @Tags Jetstream2
 */
class Benchmark {
    public static boolean debug = false;
    public static void main(String[] args) {
        runNBody();
    }

    /*
     * @Benchmark
     */
    public static double runNBody() {
        int n = 1000000;
        NBodyHandler.energy(Body.bodies);
        long start = System.nanoTime();
        for (int i = 0; i < n; i++) {
            NBodyHandler.advance(Body.bodies, 0.01);
        }
        double time = (double)(System.nanoTime() - start)/1000000;
        System.out.println("NBody - RunNBody: \t" + time + "\tms");
        return time;
    }
}
